import pickle
import os
import pandas as pd


# 保存数据
def Save(data='要保存的变量', savename='保存的名字,记得添加后缀.pkl'):
    # 保存数据文件
    data_file = open(savename , 'wb')
    pickle.dump(data, data_file)
    data_file.close()
    if os.path.exists(savename):
        print(savename+'数据保存成功!!!')
    else:
        print('数据保存失败')

# 读取数据,不存在自动读取保存数据
def Load(filename='要读取的数据文件名,记得添加后缀.pkl', ExcelName='第一次加载数据时需要输入excel地址,注意带上.xslx', SheetName='默认读取第一张表'):
    if os.path.exists(filename ):
        # 读取数据文件
        data_file = open(filename, 'rb')
        data = pickle.load(data_file)
        data_file.close()
        print(filename+'数据读取成功')
        return data
    else:
        print(filename+'数据文件不存在!!!')
        print('正在读取excel数据')
        data = 'hello world'
        if SheetName != '默认读取第一张表':
            data = pd.read_excel(ExcelName , sheet_name=SheetName)
        else:
            data = pd.read_excel(ExcelName)

        print('正在保存数据......')
        Save(data=data,savename=filename)
        return data

# 数据读取保存的整合函数
def Load_Save_Data(filename='需要保存或读取的文件名,记得添加后缀.pkl',First_time_data='第一次需要放入数据'):
    # 如果文件不存在
    if os.path.exists(filename)==False:
        # 保存文件
        Save(data=First_time_data,savename=filename)
        return First_time_data
    # 如果文件存在
    else:
        # 读取数据文件
        data_file = open(filename, 'rb')
        data = pickle.load(data_file)
        data_file.close()
        print(filename+'数据读取成功')
        return data

# 保存数据到一个excel.xlsx中的多张表
def SavetoExcel(datas='要保存的数据的列表list',sheet_name='保存到对应的表的工作表名的list列表',ExcelName='要保存的excel的名称,注意带上.xlsx'):
    if not os.path.exists(ExcelName):
        writer = pd.ExcelWriter(ExcelName)
        # 单张表
        if type(datas)!=list:
            pd.DataFrame(datas).to_excel(writer,sheet_name=sheet_name)
            writer.save()
            print(ExcelName+'保存成功,共保存了'+str(len(datas))+'行数据!!!')

        # 多张表
        else:
            for i in range(len(datas)):
                pd.DataFrame(datas[i]).to_excel(writer,sheet_name=sheet_name[i])
            # 记得保存写入流
            writer.save()
            print(ExcelName+'保存成功,共保存了'+str(len(datas))+'张表!!!')
    else:
        print(ExcelName+'已经存在,默认不重复保存')


